#include "Mesh.h"
#include "Element.h"
#include <iostream>
#include <string>
#include <cmath>
#include <fstream>
#include "MultigridSolver.h"
  

/// Eigen is linear library, providing sparse matrix, linear operator,
/// and linear solver.
#include <Eigen/Sparse>
#include <Eigen/IterativeLinearSolvers>
typedef Eigen::SparseMatrix<double> SpMat;
typedef Eigen::Triplet<double> Tri;
typedef Eigen::VectorXd RHS;
typedef Eigen::VectorXd Solution;

/** 
 * Real solution. Can be used as Dirichlet boundary conditions.
 * 
 * @param _p variable.
 * 
 * @return function value.
 */
double u(const Point<2, double> &_p);

/** 
 * Source term.
 * 
 * @param _p variable.
 * 
 * @return function value.
 */
double f(const Point<2, double> &_p);

// 编译命令
// g++ -o main test_MG_poisson.cpp -std=c++11 -I /usr/include/eigen3/ -g

int main(int argc, char* argv[])
{
	/// nx = ny = 2^n
	int n = 4; 
    //RegularMesh* mesh = new RegularMesh({0,0}, {1,1}, POW2(n), POW2(n));
	RegularMeshP2* mesh = new RegularMeshP2({0,0}, {1,1}, POW2(n), POW2(n));
    /// Need a reference element. For this example, only one kind of
    /// element.
    //P1Element<double> ref_ele;
	P2Element<double> ref_ele;

    /// The algebraic accuracy for the quadrature on the element is 1.
    ref_ele.read_quad_info("data/triangle.tmp_geo", 3);
	/// Get the quadrature information from the reference element.
	const std::valarray<Point<2, double> > &quad_pnts = ref_ele.get_quad_points();
	const std::valarray<double> &weights = ref_ele.get_quad_weights();
	int n_quad_pnts = ref_ele.get_n_quad_points();
	double volume = ref_ele.get_volume();




    size_t n_ele = mesh->get_n_grids();
	size_t n_dof = ref_ele.get_n_dofs();
    size_t n_total_dofs = mesh->get_n_dofs();
    
    /// Here we use Eigen library for all linear structures and
    /// operators.
    // Matrix<double> A(n_total_dofs, n_total_dofs);
    // std::valarray<double> rhs(0.0, n_total_dofs); / Linear algebra
    // interface. To Eigen.
    SpMat K(n_total_dofs, n_total_dofs);
    RHS Rhs(n_total_dofs);
    std::vector<Tri> TriList(n_ele * n_dof * n_dof*n_quad_pnts);
    std::vector<Tri>::iterator it = TriList.begin(); 

    /// Assembling ... For each element in the mesh.
    for (size_t ei = 0; ei < n_ele; ei++)
    {
	/// Find the structure of the current element.
	const Geometry &the_ele = mesh->get_grid(ei);
	/// Find the global indexes of all the degree of freedoms on the
	/// current element.
	const std::valarray<size_t> &dof = the_ele.get_dof();
	/// Within the indexes, we find the coordinates of the dofs
	/// from the mesh, and set to the reference element. Actually
	/// build a mapping between the global dofs and the local
	/// ones. (The indexes and coordinates of local dofs have set
	/// inside P1 element.)
	for (size_t k = 0; k < n_dof; k++)
	    ref_ele.set_global_dof(k, mesh->get_dof(dof[k]));
	
	/// Get the quadrature information from the reference element.
	const std::valarray<Point<2, double> > &quad_pnts = ref_ele.get_quad_points();
	const std::valarray<double> &weights = ref_ele.get_quad_weights();	
	int n_quad_pnts = ref_ele.get_n_quad_points();
	double volume = ref_ele.get_volume();

	/// Begin to assemble ...
	for (size_t l = 0; l < n_quad_pnts; l++)
	{
	    /// Jacobian information for the coordinates transformation.
	    double detJ = ref_ele.global2local_Jacobi_det(quad_pnts[l]);
	    Matrix<double> invJ = ref_ele.local2global_Jacobi(quad_pnts[l]);

	    /// Build the entities of the coefficent matrix and right hand side vector. 
	    for (size_t i = 0; i < n_dof; i++)
	    {
		std::valarray<double> grad_i = invJ * ref_ele.basis_gradient(i, quad_pnts[l]);
		double phi_i = ref_ele.basis_function(i, quad_pnts[l]);
		for (size_t j = 0; j < n_dof; j++)
		{
		    std::valarray<double> grad_j = invJ * ref_ele.basis_gradient(j, quad_pnts[l]);
		    /// a(grad phi_i, grad phi_j)
		    double cont = (grad_i[0] * grad_j[0] + grad_i[1] * grad_j[1]) * weights[l] * detJ * volume;
		    // A(dof[i], dof[j]) += cont;
		    *it = Tri(dof[i], dof[j], cont);
		    it++;
		}
		/// rhs(i) = \int f(x_i) phi_i dx
		// rhs[dof[i]] += f(ref_ele.local2global(quad_pnts[l])) * phi_i * weights[l] * detJ * volume;
		Rhs[dof[i]] += f(ref_ele.local2global(quad_pnts[l])) * phi_i * weights[l] * detJ * volume;
	    }
	}
    }
    K.setFromTriplets(TriList.begin(), TriList.end());
    K.makeCompressed();

//    std::cout << K << std::endl;
//    std::cout << Rhs << std::endl;
   // const std::valarray<Point<2, double> >& global_dofs = mesh->get_dofs(); 

    for (auto bndIdx : mesh->get_all_boundary())
    {
		Point<2, double> bndDof = mesh->get_dof(bndIdx);
	    double boundary_value = u(bndDof);
	    Rhs[bndIdx] = boundary_value * K.coeffRef(bndIdx, bndIdx);
	    for (Eigen::SparseMatrix<double>::InnerIterator it(K, bndIdx); it; ++it)
	    {
		size_t row = it.row();
		if (row == bndIdx)
		    continue;
		Rhs[row] -= K.coeffRef(bndIdx, row) * boundary_value;
		K.coeffRef(bndIdx, row) = 0.0;
		K.coeffRef(row, bndIdx) = 0.0;
	    }
    }; 

    
	Eigen::ConjugateGradient<Eigen::SparseMatrix<double> > Solver_sparse;
    Eigen::VectorXd x1_sparse;
    Solver_sparse.setTolerance(1e-15);
    Solver_sparse.compute(K);
    Solution solution1 = Solver_sparse.solve(Rhs);
	//P1_MGSolver<SpMat> Solver(K, Rhs, u, mesh, n);
	P2_MGSolver<SpMat> Solver(K, Rhs, u, mesh, n);
	Solver.compute(K);
    Solver.setTolerance(1e-12);
    VectorXd solution = Solver.Solve();
	/*
	double diff = 0;
	for (int i = 0; i < solution1.size(); i++)
		diff += (solution1(i) - solution1(i)) * (solution(i) - solution(i));
	std::cout << "和CG求解器的偏差: " << sqrt(diff) << std::endl; */
};

double u(const Point<2, double> &_p)
{
    return std::sin(M_PI * _p[0]) * std::sin(M_PI * _p[1]);
};

double f(const Point<2, double> &_p)
{
    return 2 * M_PI * M_PI * u(_p);
};

/*double u(const Point<2, double> &_p)
{
    return std::cos(M_PI * _p[0]) * std::cos(M_PI * _p[1]);
};

double f(const Point<2, double> &_p)
{
    return 2 * M_PI * M_PI * u(_p);
};*/

